package com.luo.util;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.jsoup.Jsoup;
import org.jsoup.safety.Whitelist;
import org.springframework.core.io.ClassPathResource;

import java.io.IOException;
import java.io.InputStream;
import java.util.Iterator;



/**
 * XSS工具
 *
 * @author withwindluo
 * @since 2021/12/9 18:52
 */
public class XSSUtil {

    private static final ClassPathResource jsoupWhiteListPathRes = new ClassPathResource("/json/xssWhiteList.json");

    //添加默认base配置,因为本项目的富文本图片实现使用base64,暂不使用默认的baseImage配置
//    private static Whitelist whitelist = Whitelist.basicWithImages();
    private static Whitelist whitelist = Whitelist.basic();

    //再载入json自定义白名单
    static {
        InputStream whiteConfig = null;
        try {
            whiteConfig = jsoupWhiteListPathRes.getInputStream();
        } catch (IOException e) {
            e.printStackTrace();
        }

        if (whiteConfig == null) {
            throw new RuntimeException("读取jsoup xss 白名单文件失败");
        } else {
            try {
                JSONObject whiteListJson = JSON.parseObject(whiteConfig, JSONObject.class);

                //添加标签 addTags
                JSONArray addTagsJsonArr = whiteListJson.getJSONArray("addTags");
                String[] addTagsArr = addTagsJsonArr.toArray(new String[0]);
                whitelist.addTags(addTagsArr);


                //添加属性 addAttributes
                JSONArray addAttrJsonArr = whiteListJson.getJSONArray("addAttributes");
                Iterator<Object> iter = addAttrJsonArr.iterator();
                while (iter.hasNext()) {
                    JSONObject attrJsonObj = (JSONObject) iter.next();
                    String tag = attrJsonObj.getString("tag");
                    JSONArray attrJsonArr = attrJsonObj.getJSONArray("attributes");
                    String[] attrArr = attrJsonArr.toArray(new String[0]);
                    whitelist.addAttributes(tag, attrArr);
                }


                //添加 addProtocols
                JSONArray addProtoJsonArr = whiteListJson.getJSONArray("addProtocols");
                iter = addProtoJsonArr.iterator();
                while (iter.hasNext()) {
                    JSONObject attrJsonObj = (JSONObject) iter.next();
                    String tag = attrJsonObj.getString("tag");
                    String attribute = attrJsonObj.getString("attribute");
                    JSONArray protoJsonArr = attrJsonObj.getJSONArray("protocols");
                    String[] protocolArr = protoJsonArr.toArray(new String[0]);
                    whitelist.addProtocols(tag, attribute, protocolArr);
                }


            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private XSSUtil() {

    }

    /**
     * 使用jsoup设置标签放行白名单
     *
     * @param originStr
     * @return
     */
    public static String jsoupCleanRichText(String originStr) {

//        whitelist = (new Whitelist()
//                .addTags("a", "b", "div", "img", "p", "strong") // 设置允许的标签
//                .addAttributes("a", "href", "title", "...") // 设置标签允许的属性
//                .addAttributes(":all", "class", "id", "src") // 通配符，对所有标签配置允许的属性
//                .addProtocols("img", "src", "http", "https")); // 设置Protocol，这是代表img的src属性只允许http和https开头

        return Jsoup.clean(originStr, whitelist);
    }
}
